"""Usage tracking system for API keys."""

import asyncio
from typing import Dict, Any, Optional, List
from datetime import datetime, timedelta
from dataclasses import dataclass, field
import motor.motor_asyncio
from bson import ObjectId
from fastapi import APIRouter, HTTPException, Depends, Query
from pydantic import BaseModel
import redis.asyncio as redis

from ..models.api_key_model import APIKey
from ..api.key_management import get_db, get_current_user


router = APIRouter(prefix="/api/keys", tags=["Usage Tracking"])


@dataclass
class UsageRecord:
    """Usage record for tracking API key usage."""
    key_id: str
    timestamp: datetime = field(default_factory=datetime.utcnow)
    tokens_used: int = 0
    cost: float = 0.0
    model: str = ""
    success: bool = True
    request_id: Optional[str] = None
    error: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for MongoDB storage."""
        return {
            "key_id": self.key_id,
            "timestamp": self.timestamp,
            "tokens_used": self.tokens_used,
            "cost": self.cost,
            "model": self.model,
            "success": self.success,
            "request_id": self.request_id,
            "error": self.error
        }


class UsageStats(BaseModel):
    """Usage statistics response model."""
    total_tokens: int
    total_cost: float
    request_count: int
    success_rate: float
    period: str
    breakdown_by_model: Dict[str, Dict[str, Any]]
    daily_usage: List[Dict[str, Any]]


class UsageTracker:
    """Tracks API key usage and costs."""
    
    def __init__(self, mongodb_client: motor.motor_asyncio.AsyncIOMotorClient = None,
                 redis_client: redis.Redis = None):
        """Initialize usage tracker.
        
        Args:
            mongodb_client: MongoDB client for persistent storage
            redis_client: Redis client for real-time tracking
        """
        self.mongodb = mongodb_client
        self.redis = redis_client
        self._usage_cache: Dict[str, List[UsageRecord]] = {}
        self._flush_task = None
        
    async def start(self):
        """Start the usage tracker background tasks."""
        if not self._flush_task:
            self._flush_task = asyncio.create_task(self._flush_loop())
            
    async def stop(self):
        """Stop the usage tracker."""
        if self._flush_task:
            self._flush_task.cancel()
            await self._flush_to_db()
            
    async def track_usage(
        self,
        key_id: str,
        tokens_used: int,
        cost: float,
        model: str,
        success: bool = True,
        request_id: Optional[str] = None,
        error: Optional[str] = None
    ):
        """Track usage of an API key.
        
        Args:
            key_id: API key ID
            tokens_used: Number of tokens consumed
            cost: Estimated cost in USD
            model: Model name used
            success: Whether the request was successful
            request_id: Optional request ID for tracing
            error: Error message if failed
        """
        record = UsageRecord(
            key_id=key_id,
            tokens_used=tokens_used,
            cost=cost,
            model=model,
            success=success,
            request_id=request_id,
            error=error
        )
        
        # Add to cache
        if key_id not in self._usage_cache:
            self._usage_cache[key_id] = []
        self._usage_cache[key_id].append(record)
        
        # Update Redis cache for real-time stats
        if self.redis:
            await self._update_redis_stats(key_id, record)
            
        # Update MongoDB immediately for important records
        if not success or cost > 1.0:  # Failed requests or high-cost operations
            await self._write_to_db([record])
            
    async def _update_redis_stats(self, key_id: str, record: UsageRecord):
        """Update real-time statistics in Redis with fallback to MongoDB.
        
        Args:
            key_id: API key ID
            record: Usage record
        """
        if not self.redis:
            await self._fallback_to_mongodb_stats(key_id, record)
            return
            
        try:
            # Update daily counters
            today = datetime.utcnow().strftime("%Y-%m-%d")
            daily_key = f"usage:{key_id}:{today}"
            
            pipe = self.redis.pipeline()
            pipe.hincrby(daily_key, "tokens", record.tokens_used)
            pipe.hincrbyfloat(daily_key, "cost", record.cost)
            pipe.hincrby(daily_key, "requests", 1)
            if record.success:
                pipe.hincrby(daily_key, "success", 1)
            pipe.expire(daily_key, 86400 * 7)  # Keep for 7 days
            
            # Update monthly counters
            month = datetime.utcnow().strftime("%Y-%m")
            monthly_key = f"usage:{key_id}:{month}"
            pipe.hincrby(monthly_key, "tokens", record.tokens_used)
            pipe.hincrbyfloat(monthly_key, "cost", record.cost)
            pipe.hincrby(monthly_key, "requests", 1)
            pipe.expire(monthly_key, 86400 * 35)  # Keep for 35 days
            
            await pipe.execute()
        except (redis.ConnectionError, redis.TimeoutError, Exception) as e:
            print(f"Redis update failed, falling back to MongoDB: {e}")
            await self._fallback_to_mongodb_stats(key_id, record)
            
    async def _flush_loop(self):
        """Background loop to flush usage data to database."""
        while True:
            try:
                await asyncio.sleep(60)  # Flush every minute
                await self._flush_to_db()
            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"Flush error: {e}")
                
    async def _flush_to_db(self):
        """Flush cached usage records to database."""
        if not self._usage_cache or not self.mongodb:
            return
            
        # Copy and clear cache
        records_to_write = []
        key_updates = {}
        
        for key_id, records in self._usage_cache.items():
            records_to_write.extend(records)
            
            # Aggregate stats for key update
            total_tokens = sum(r.tokens_used for r in records)
            total_cost = sum(r.cost for r in records)
            request_count = len(records)
            
            key_updates[key_id] = {
                "tokens": total_tokens,
                "cost": total_cost,
                "requests": request_count
            }
            
        self._usage_cache.clear()
        
        # Write to database
        if records_to_write:
            await self._write_to_db(records_to_write)
            
        # Update API key stats
        for key_id, stats in key_updates.items():
            await self._update_key_stats(key_id, stats)
            
    async def _write_to_db(self, records: List[UsageRecord]):
        """Write usage records to database.
        
        Args:
            records: List of usage records to write
        """
        if not self.mongodb:
            return
            
        try:
            db = self.mongodb.ai_writing
            docs = [r.to_dict() for r in records]
            await db.usage_records.insert_many(docs)
        except Exception as e:
            print(f"Database write failed: {e}")
            
    async def _update_key_stats(self, key_id: str, stats: Dict[str, Any]):
        """Update API key usage statistics.
        
        Args:
            key_id: API key ID
            stats: Statistics to update
        """
        if not self.mongodb:
            return
            
        try:
            db = self.mongodb.ai_writing
            await db.api_keys.update_one(
                {"_id": ObjectId(key_id)},
                {
                    "$inc": {
                        "usage_stats.total_tokens": stats["tokens"],
                        "usage_stats.total_cost": stats["cost"],
                        "usage_stats.request_count": stats["requests"]
                    },
                    "$set": {
                        "last_used_at": datetime.utcnow()
                    }
                }
            )
        except Exception as e:
            print(f"Key stats update failed: {e}")
            
    async def _fallback_to_mongodb_stats(self, key_id: str, record: UsageRecord):
        """Fallback method to update stats directly in MongoDB when Redis fails.
        
        Args:
            key_id: API key ID
            record: Usage record
        """
        if not self.mongodb:
            print("Warning: Both Redis and MongoDB unavailable for stats update")
            return
            
        try:
            db = self.mongodb.ai_writing
            
            # Create aggregated daily stats document in MongoDB
            today = datetime.utcnow().strftime("%Y-%m-%d")
            daily_stats_key = f"daily_stats_{key_id}_{today}"
            
            # Use upsert to create or update daily stats
            await db.usage_stats.update_one(
                {"key_id": key_id, "date": today, "type": "daily"},
                {
                    "$inc": {
                        "total_tokens": record.tokens_used,
                        "total_cost": record.cost,
                        "request_count": 1,
                        "success_count": 1 if record.success else 0
                    },
                    "$set": {
                        "last_updated": datetime.utcnow()
                    }
                },
                upsert=True
            )
            
            # Also update monthly stats
            month = datetime.utcnow().strftime("%Y-%m")
            await db.usage_stats.update_one(
                {"key_id": key_id, "date": month, "type": "monthly"},
                {
                    "$inc": {
                        "total_tokens": record.tokens_used,
                        "total_cost": record.cost,
                        "request_count": 1,
                        "success_count": 1 if record.success else 0
                    },
                    "$set": {
                        "last_updated": datetime.utcnow()
                    }
                },
                upsert=True
            )
            
        except Exception as e:
            print(f"MongoDB fallback stats update failed: {e}")
            
    async def get_usage_stats(
        self,
        key_id: str,
        period: str = "today",
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None
    ) -> UsageStats:
        """Get usage statistics for an API key.
        
        Args:
            key_id: API key ID
            period: Time period (today, week, month, custom)
            start_date: Start date for custom period
            end_date: End date for custom period
            
        Returns:
            Usage statistics
        """
        # Determine date range
        if period == "today":
            start = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0)
            end = datetime.utcnow()
        elif period == "week":
            start = datetime.utcnow() - timedelta(days=7)
            end = datetime.utcnow()
        elif period == "month":
            start = datetime.utcnow() - timedelta(days=30)
            end = datetime.utcnow()
        else:  # custom
            start = start_date or datetime.utcnow() - timedelta(days=30)
            end = end_date or datetime.utcnow()
            
        # Try Redis first for recent data
        if self.redis and period in ["today", "week"]:
            stats = await self._get_redis_stats(key_id, start, end)
            if stats:
                return stats
                
        # Fall back to MongoDB
        if self.mongodb:
            return await self._get_db_stats(key_id, start, end)
            
        return UsageStats(
            total_tokens=0,
            total_cost=0.0,
            request_count=0,
            success_rate=0.0,
            period=period,
            breakdown_by_model={},
            daily_usage=[]
        )
        
    async def _get_redis_stats(
        self,
        key_id: str,
        start: datetime,
        end: datetime
    ) -> Optional[UsageStats]:
        """Get usage statistics from Redis.
        
        Args:
            key_id: API key ID
            start: Start date
            end: End date
            
        Returns:
            Usage statistics or None if not available
        """
        try:
            total_tokens = 0
            total_cost = 0.0
            total_requests = 0
            total_success = 0
            daily_usage = []
            
            current = start
            while current <= end:
                day_key = f"usage:{key_id}:{current.strftime('%Y-%m-%d')}"
                data = await self.redis.hgetall(day_key)
                
                if data:
                    day_tokens = int(data.get(b"tokens", 0))
                    day_cost = float(data.get(b"cost", 0))
                    day_requests = int(data.get(b"requests", 0))
                    day_success = int(data.get(b"success", 0))
                    
                    total_tokens += day_tokens
                    total_cost += day_cost
                    total_requests += day_requests
                    total_success += day_success
                    
                    daily_usage.append({
                        "date": current.strftime("%Y-%m-%d"),
                        "tokens": day_tokens,
                        "cost": day_cost,
                        "requests": day_requests
                    })
                    
                current += timedelta(days=1)
                
            success_rate = (total_success / total_requests * 100) if total_requests > 0 else 0
            
            return UsageStats(
                total_tokens=total_tokens,
                total_cost=total_cost,
                request_count=total_requests,
                success_rate=success_rate,
                period=f"{start.date()} to {end.date()}",
                breakdown_by_model={},  # Not available in Redis
                daily_usage=daily_usage
            )
        except Exception:
            return None
            
    async def _get_db_stats(
        self,
        key_id: str,
        start: datetime,
        end: datetime
    ) -> UsageStats:
        """Get usage statistics from MongoDB.
        
        Args:
            key_id: API key ID
            start: Start date
            end: End date
            
        Returns:
            Usage statistics
        """
        db = self.mongodb.ai_writing
        
        # Aggregation pipeline
        pipeline = [
            {
                "$match": {
                    "key_id": key_id,
                    "timestamp": {"$gte": start, "$lte": end}
                }
            },
            {
                "$group": {
                    "_id": None,
                    "total_tokens": {"$sum": "$tokens_used"},
                    "total_cost": {"$sum": "$cost"},
                    "request_count": {"$sum": 1},
                    "success_count": {
                        "$sum": {"$cond": ["$success", 1, 0]}
                    }
                }
            }
        ]
        
        # Get overall stats
        cursor = db.usage_records.aggregate(pipeline)
        results = await cursor.to_list(length=1)
        
        if results:
            stats = results[0]
            total_tokens = stats["total_tokens"]
            total_cost = stats["total_cost"]
            request_count = stats["request_count"]
            success_rate = (stats["success_count"] / request_count * 100) if request_count > 0 else 0
        else:
            total_tokens = 0
            total_cost = 0.0
            request_count = 0
            success_rate = 0.0
            
        # Get breakdown by model
        model_pipeline = [
            {
                "$match": {
                    "key_id": key_id,
                    "timestamp": {"$gte": start, "$lte": end}
                }
            },
            {
                "$group": {
                    "_id": "$model",
                    "tokens": {"$sum": "$tokens_used"},
                    "cost": {"$sum": "$cost"},
                    "requests": {"$sum": 1}
                }
            }
        ]
        
        model_cursor = db.usage_records.aggregate(model_pipeline)
        model_results = await model_cursor.to_list(length=None)
        
        breakdown_by_model = {
            result["_id"]: {
                "tokens": result["tokens"],
                "cost": result["cost"],
                "requests": result["requests"]
            }
            for result in model_results
        }
        
        # Get daily usage
        daily_pipeline = [
            {
                "$match": {
                    "key_id": key_id,
                    "timestamp": {"$gte": start, "$lte": end}
                }
            },
            {
                "$group": {
                    "_id": {
                        "$dateToString": {
                            "format": "%Y-%m-%d",
                            "date": "$timestamp"
                        }
                    },
                    "tokens": {"$sum": "$tokens_used"},
                    "cost": {"$sum": "$cost"},
                    "requests": {"$sum": 1}
                }
            },
            {
                "$sort": {"_id": 1}
            }
        ]
        
        daily_cursor = db.usage_records.aggregate(daily_pipeline)
        daily_results = await daily_cursor.to_list(length=None)
        
        daily_usage = [
            {
                "date": result["_id"],
                "tokens": result["tokens"],
                "cost": result["cost"],
                "requests": result["requests"]
            }
            for result in daily_results
        ]
        
        return UsageStats(
            total_tokens=total_tokens,
            total_cost=total_cost,
            request_count=request_count,
            success_rate=success_rate,
            period=f"{start.date()} to {end.date()}",
            breakdown_by_model=breakdown_by_model,
            daily_usage=daily_usage
        )


# Global tracker instance
tracker = UsageTracker()


@router.get("/{key_id}/usage", response_model=UsageStats)
async def get_key_usage(
    key_id: str,
    period: str = Query("today", description="Time period: today, week, month, custom"),
    start_date: Optional[datetime] = Query(None, description="Start date for custom period"),
    end_date: Optional[datetime] = Query(None, description="End date for custom period"),
    current_user: Dict = Depends(get_current_user),
    db=Depends(get_db)
):
    """Get usage statistics for an API key."""
    # Verify ownership
    doc = await db.api_keys.find_one({
        "_id": ObjectId(key_id),
        "user_id": current_user["id"]
    })
    
    if not doc:
        raise HTTPException(status_code=404, detail="API key not found")
    
    # Initialize tracker with database
    tracker.mongodb = db.client
    
    # Get usage stats
    stats = await tracker.get_usage_stats(
        key_id=key_id,
        period=period,
        start_date=start_date,
        end_date=end_date
    )
    
    return stats


# Middleware to track usage
class UsageTrackingMiddleware:
    """Middleware to track API key usage in model calls."""
    
    def __init__(self, tracker: UsageTracker):
        """Initialize middleware.
        
        Args:
            tracker: Usage tracker instance
        """
        self.tracker = tracker
        
    async def __call__(self, key_id: str, model: str, tokens: int, cost: float):
        """Track usage for a model call.
        
        Args:
            key_id: API key ID
            model: Model name
            tokens: Tokens used
            cost: Estimated cost
        """
        await self.tracker.track_usage(
            key_id=key_id,
            tokens_used=tokens,
            cost=cost,
            model=model,
            success=True
        )